{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "8.2_async-computation.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W3EYDWKmn16n",
        "colab_type": "text"
      },
      "source": [
        "# 8.2 异步计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lyqEojzxnlxh",
        "colab_type": "code",
        "outputId": "d262a478-4cf8-466c-ef05-2719ac4edd09",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "# 安装 TensorFlow\n",
        "try:\n",
        "  # Colab only\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "    pass"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 2.x selected.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lp3Nz4GTn8Wc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.keras as keras\n",
        "import os\n",
        "import subprocess\n",
        "import time"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-jhEFI9nz9P",
        "colab_type": "text"
      },
      "source": [
        "## 8.2.1 Tensorflow 中的异步计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "th7xxlN0oUJ6",
        "colab_type": "code",
        "outputId": "3ba908c2-0081-4ed6-94f2-503dd9684978",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "a = tf.ones((1, 2))\n",
        "b = tf.ones((1, 2))\n",
        "c = a * b + 2\n",
        "c"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[3., 3.]], dtype=float32)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S4PXsRZWoj33",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Benchmark(object):\n",
        "  def __init__(self, prefix=None):\n",
        "    self.prefix = prefix + ' ' if prefix else ''\n",
        "\n",
        "  def __enter__(self):\n",
        "    self.start = time.time()\n",
        "\n",
        "  def __exit__(self, *args):\n",
        "    print('%stime: %.4f sec' % (self.prefix, time.time() - self.start))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pXbvbaHlsNV4",
        "colab_type": "code",
        "outputId": "0736f9c4-e3a4-41f9-8d0a-171e149ccf70",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "with Benchmark('Workloads are queued.'):\n",
        "  x = tf.random.uniform(shape=(2000, 2000))\n",
        "  y = tf.keras.backend.sum(tf.transpose(x) * x)\n",
        "\n",
        "with Benchmark('Workloads are finished.'):\n",
        "  print('sum =', y)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Workloads are queued. time: 0.1726 sec\n",
            "sum = tf.Tensor(1000411.9, shape=(), dtype=float32)\n",
            "Workloads are finished. time: 0.0006 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q_jPu3SpuX4i",
        "colab_type": "text"
      },
      "source": [
        "## 8.2.2 用同步函数让前端等待计算结果"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZqRhEbxvsgKG",
        "colab_type": "code",
        "outputId": "986eb49e-1d95-4299-bce5-108d8eed8153",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "with Benchmark():\n",
        "  y = tf.keras.backend.sum(tf.transpose(x) * x)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.0341 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BGd6zqszvQSQ",
        "colab_type": "code",
        "outputId": "3af1d420-0ecd-404e-cf50-4e2759cf8844",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "with Benchmark():\n",
        "  y = tf.keras.backend.sum(tf.transpose(x) * x)\n",
        "  z = tf.keras.backend.sum(tf.transpose(x) * x)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.0386 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fhqKiTjuvYJI",
        "colab_type": "code",
        "outputId": "ddeaabd9-d089-42ce-fdad-7c9acfeabca4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "with Benchmark():\n",
        "  y = tf.keras.backend.sum(tf.transpose(x) * x)\n",
        "  y.numpy()"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.0223 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dchOO-OPvpFB",
        "colab_type": "code",
        "outputId": "e4ba3a4c-ad17-44c7-9d4f-293f82b6e600",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "with Benchmark():\n",
        "  y = tf.keras.backend.sum(tf.transpose(x) * x)\n",
        "  print(tf.norm(y).numpy())"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1000411.9\n",
            "time: 0.0463 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WWFDBnMwx18_",
        "colab_type": "text"
      },
      "source": [
        "## 8.2.3 使用异步计算提升计算性能"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DR8hJqgUxoZN",
        "colab_type": "code",
        "outputId": "f07cd855-96b3-45c6-a0d9-b838b68e7398",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "with Benchmark('synchronous.'):\n",
        "  for _ in range(1000):\n",
        "    y = x + 1\n",
        "\n",
        "@tf.function\n",
        "def loop():\n",
        "  for _ in range(1000):\n",
        "    y = x + 1\n",
        "  return y\n",
        "\n",
        "with Benchmark('asynchronous.'):\n",
        "  y = loop()"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "synchronous. time: 3.1374 sec\n",
            "asynchronous. time: 1.1313 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4N3fOI523RUh",
        "colab_type": "text"
      },
      "source": [
        "## 8.2.4 异步计算对内存的影响"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dj-X-x2K1LBU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def data_iter():\n",
        "  start = time.time()\n",
        "  num_batches, batch_size = 100, 1024\n",
        "  for i in range(num_batches):\n",
        "    X = tf.random.normal(shape=(batch_size, 512))\n",
        "    y = tf.ones((batch_size,))\n",
        "    yield X, y\n",
        "    if (i + 1) % 50 == 0:\n",
        "      print('batch %d, time %f sec' % (i+1, time.time()-start))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pr8mN2mB4RmF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "net = keras.Sequential()\n",
        "net.add(keras.layers.Dense(2048, activation='relu'))\n",
        "net.add(keras.layers.Dense(512, activation='relu'))\n",
        "net.add(keras.layers.Dense(1))\n",
        "optimizer=keras.optimizers.SGD(0.05)\n",
        "loss = keras.losses.MeanSquaredError()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WGnntfFLmHJZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_mem():\n",
        "  res = subprocess.check_output(['ps', 'u', '-p', str(os.getpid())])\n",
        "  return int(str(res).split()[15]) / 1e3"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RCAcJBp6nkWl",
        "colab_type": "code",
        "outputId": "2d5699d4-681c-4065-f9d4-c6b36543f3dd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "for X, y in data_iter():\n",
        "  break\n",
        "loss(y, net(X))"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor: shape=(), dtype=float32, numpy=0.49068463>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4cuicIm_Jc9q",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "81017d8b-8e49-4f1c-d4eb-ca07b6239866"
      },
      "source": [
        "l_sum, mem = 0, get_mem()\n",
        "dense_1 = keras.layers.Dense(2048, activation='relu')\n",
        "dense_2 = keras.layers.Dense(512, activation='relu')\n",
        "dense_3 = keras.layers.Dense(1)\n",
        "trainable_variables = (dense_1.trainable_variables + \n",
        "                       dense_2.trainable_variables +\n",
        "                       dense_3.trainable_variables)\n",
        "for X, y in data_iter():\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = net(X)\n",
        "    loss_value = loss(y, logits)\n",
        "\n",
        "  grads = tape.gradient(loss_value, trainable_variables)\n",
        "  optimizer.apply_gradients(zip(grads, trainable_variables))\n",
        "\n",
        "print('increased memory: %f MB' % (get_mem() - mem))"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "batch 50, time 7.880550 sec\n",
            "batch 100, time 15.700529 sec\n",
            "increased memory: 14.336000 MB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-agmYI7ZbBkY",
        "colab_type": "code",
        "outputId": "5a09bf78-3245-49f5-d036-4c8b770e1b8f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "l_sum, mem = 0, get_mem()\n",
        "for X, y in data_iter():\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = net(X)\n",
        "    loss_value = loss(y, logits)\n",
        "\n",
        "  grads = tape.gradient(loss_value, net.trainable_weights)\n",
        "  optimizer.apply_gradients(zip(grads, net.trainable_weights))\n",
        "\n",
        "print('increased memory: %f MB' % (get_mem() - mem))"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "batch 50, time 7.976524 sec\n",
            "batch 100, time 15.683179 sec\n",
            "increased memory: 12.268000 MB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vqS9hgNcitGQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}